import torch


class STFTEncoder:
    """Pure STFT encoder"""

    def __init__(
        self,
        n_fft=256,
        win_length=256,
        hop_length=64,
        window="hann",
        center=True,
        normalized=False,
        onesided=True,
        return_complex=True,
        device="cuda",
    ):
        window_func = getattr(torch, f"{window}_window")
        self.window_tensor = window_func(n_fft).to(device)
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.win_length = win_length
        self.center = center
        self.normalized = normalized
        self.onesided = onesided
        self.return_complex = return_complex
        pass

    def __call__(self, input: torch.Tensor):
        """
        input shape: [B, 1, T (32000)]
        Apply stft on the input signal
        output shape: [B, 2 (RL), Time, Feature (129)]
        """
        input = input.squeeze(1)
        x = torch.stft(
            input,
            n_fft=self.n_fft,
            win_length=self.win_length,
            hop_length=self.hop_length,
            window=self.window_tensor,
            center=self.center,
            normalized=self.normalized,
            onesided=self.onesided,
            return_complex=self.return_complex,
        )
        x = torch.view_as_real(x)
        x = x.permute(0, 3, 2, 1)
        return x


class STFTDecoder:
    def __init__(
        self,
        n_fft=256,
        win_length=256,
        hop_length=64,
        window="hann",
        center=True,
        normalized=False,
        onesided=True,
        return_complex=False,
        src=2,
        device="cuda",
    ):
        window_func = getattr(torch, f"{window}_window")
        self.window_tensor = window_func(n_fft).to(device)
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.win_length = win_length
        self.center = center
        self.normalized = normalized
        self.onesided = onesided
        self.return_complex = return_complex
        self.src = src
        pass

    def __call__(self, input: torch.Tensor, b):
        """
        Input shape: [B*src, T, F]
        """
        # [B*src, F, T]

        x = input.transpose(1, 2)

        output = torch.istft(
            x,
            n_fft=self.n_fft,
            hop_length=self.hop_length,
            win_length=self.win_length,
            window=self.window_tensor,
            center=self.center,
            normalized=self.normalized,
            onesided=self.onesided,
            return_complex=self.return_complex,
        )
        output = output.view(b, self.src, -1)

        return output

    pass


if __name__ == "__main__":
    s = STFTEncoder()
    x = torch.randn(3, 1, 32000)
    y = s(x)
    print(y.shape)
